import os

os.environ['WANDB_DISABLED'] = 'true'
from ultralytics import YOLO
from ultralytics.data import build_dataloader, build_llltData
from ultralytics.cfg import get_cfg
from ultralytics.models.yolo.detect import DetectionTrainer
# 添加运行参数
from argparse import ArgumentParser
# yaml解析
import yaml
import torch.distributed as dist
import torch
from ultralytics.utils.torch_utils import select_device

parser = ArgumentParser()
parser.add_argument("--data", type=str, default="coco.yaml", help="data yaml file")
parser.add_argument("--model", type=str, default="snn_yolov8l.yaml", help="model yaml file")
parser.add_argument("--epochs", type=int, default=100, help="epochs num")
parser.add_argument("--batch", type=int, default=24, help="batch size")
parser.add_argument("--workers", type=int, default=4, help="workers num")
# 获取参数


args = parser.parse_args()
with open("./ultralytics/cfg/datasets/"+ args.data, "r") as file:
    data = yaml.safe_load(file)
image_path = data.get('path')
device = [0]


model =YOLO("snn_yolov8l.yaml")

# print(model)
#train coco
# model.train(data="coco.yaml",device=[7],epochs=100)  # train the model
model.train(data="coco.yaml",device=device,epochs=100, resume=True, batch=args.batch)  # train the model

#TEST
# model = YOLO('runs/detect/train1/weights/last.pt')  # load a pretrained model (recommended for training)

